from gensim.models import Word2Vec
from gensim.models import KeyedVectors
import os
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# 使用 Word2Vec 训练自定义语料
# 示例语料
sentences = [  
    ["猫", "喜欢", "晒太阳"],  
    ["狗", "猫","是", "人类", "最好的", "朋友"],  
    ["猫", "和", "狗", "是", "常见", "的", "宠物"],  
    ["宠物", "可以", "给", "人", "带来", "快乐"]]

# 训练 Word2Vec 模型
word2vec_model = Word2Vec(sentences, vector_size=50, window=3, min_count=1, workers=4)

# 打印词汇的嵌入向量
print("Word2Vec 生成的词向量: ")
print("猫:", word2vec_model.wv["猫"])
print("狗:", word2vec_model.wv["狗"])

# 计算词语相似性
similarity = word2vec_model.wv.similarity("猫", "狗")
print(f"\nWord2Vec 中 '猫' 和 '狗' 的相似性: {similarity:.4f}")

# 加载预训练的 GloVe 向量
# 假设已下载 GloVe 文件 glove.6B.50d.txt 并解压
glove_path = "vector-base-glove.6B.50d.txt"

# 将 GloVe 转换为 Gensim 格式
def load_glove_model(glove_file):
    glove_vectors = {}
    with open(glove_file, "r", encoding="utf-8") as f:
        for line in f:
            split_line = line.split()
            word = split_line[0]
            embedding = list(map(float, split_line[1:]))
            glove_vectors[word] = embedding
    return glove_vectors

if os.path.exists(glove_path):
    glove_model = load_glove_model(glove_path)

# 查询 GloVe 向量
word = "dog"
if word in glove_model:
    print(f"\nGloVe 中 '{word}' 的向量:")
    print(glove_model[word][:15], "...")  # 只展示部分维度

# 使用 GloVe 和 Word2Vec 的对比分析
# 比较两者的相似性（假设 GloVe加载完成）
if "cat" in glove_model and "dog" in glove_model:
    glove_similarity = cosine_similarity([glove_model["cat"]], [glove_model["dog"]])[0][0]
    print(f"\nGloVe 中 'cat' 和 'dog' 的相似性: {glove_similarity:.4f}")